import torch
from PIL import Image
from torchvision.transforms import ToTensor
from model import CustomNet as MyModel  # 从你的model.py导入模型类


def process_image(image_path):
    """处理单张图像并返回识别结果"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载模型（新版state_dict方式）
    model = MyModel().to(device)  # 实例化你的模型结构
    model.load_state_dict(torch.load('models/final_model.pth', map_location=device))
    model.eval()

    # 图像预处理
    image = Image.open(image_path)
    transform = ToTensor()
    image_tensor = transform(image).unsqueeze(0).to(device)

    # 推理
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output.data, 1)
        class_id = predicted.item()

    # 示例标签（根据你的实际类别修改）
    labels = {
        0: "拳头",
        1: "手掌",
        2: "剪刀手",
        3: "OK手势",
        4: "竖起大拇指"
    }

    return labels.get(class_id, "未知手势")